from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.naive_bayes import MultinomialNB
from sklearn.neighbors import KNeighborsClassifier
from sklearn.ensemble import RandomForestClassifier
from dataPretreat import data_pretreat


def mx_line(x_train, y_train):
    # 线性回归算法
    mx = LinearRegression()
    mx.fit(x_train, y_train)
    return mx


def mx_log(x_train, y_train):
    # 逻辑回归算法
    mx = LogisticRegression(penalty=12)
    mx.fit(x_train, y_train)
    return mx


def mx_bayes(x_train, y_train):
    # 多项式朴素贝叶斯算法
    mx = MultinomialNB(alpha=0.01)
    mx.fit(x_train, y_train)
    return mx


def mx_knn(x_train, y_train):
    # KNN邻近算法
    mx = KNeighborsClassifier()
    mx.fit(x_train, y_train)
    return mx


def mx_forest(x_train, y_train):
    # 随机森林算法
    mx = RandomForestClassifier(n_estimators=8)
    mx.fit(x_train, y_train)
    return mx


if __name__ == '__main__':
    data_train = data_pretreat(path='../resource/KDDTrain+.csv')
    # 根据论文选择的特征值
    """
    https://mathpretty.com/10244.html#Misuse_detection
    """
    x = data_train[['duration', 'protocol_type', 'service', 'src_bytes', 'dst_bytes', 'wrong_fragment', 'serror_rate',
                    'dst_host_srv_count', 'dst_host_diff_srv_rate', 'dst_host_same_src_port_rate',
                    'dst_host_rerror_rate']]
    y = data_train['label']
    # x_train, x_test, y_train, y_test = train_test_split(x, y, random_state=2)
    # 生成模型mx
    mx = mx_forest(x.values, y.values)
    # 用模型mx预测x_test里的数据并返回结果集y_pred
    data_test = data_pretreat(path='../resource/KDDTest+.csv')
    x_test = data_test[
        ['duration', 'protocol_type', 'service', 'src_bytes', 'dst_bytes', 'wrong_fragment', 'serror_rate',
         'dst_host_srv_count', 'dst_host_diff_srv_rate', 'dst_host_same_src_port_rate',
         'dst_host_rerror_rate']]
    y_test = data_test['label']
    y_pred = mx.predict(x_test.values)
    # print(y_pred)
    # print(y_test)
    '''
    模型评估
    参考网站:https://blog.csdn.net/sinat_16388393/article/details/91427631
    '''
    print("准确率为:" + str(mx.score(x_test, y_test)))
    '''target_names = ['normal', 'ipsweep', 'mscan', 'nmap', 'portsweep', 'saint', 'satan', 'apache2',
                    'back', 'land', 'mailbomb', 'neptune', 'pod', 'processtable', 'smurf',
                    'teardrop', 'udpstorm', 'buffer_overflow', 'httptunnel', 'loadmodule', 'perl',
                    'ps', 'rootkit', 'sqlattack', 'xterm', 'ftp_write', 'guess_passwd', 'imap',
                    'multihop', 'named', 'phf', 'sendmail', 'snmpgetattack', 'snmpguess', 'spy',
                    'warezclient', 'warezmaster', 'worm', 'xlock', 'xsnoop']
    labels = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
              30, 31, 32, 33, 34, 35, 36, 37, 38, 39]
    print(classification_report(y_test, y_pred, labels=labels, target_names=target_names))
    # 混淆矩阵
    confusion_mat = confusion_matrix(y_test, y_pred)
    print(confusion_mat)
    # ROC AUC 值
    # roc_auc_score(y_test, y_pred)'''
